import os
import argparse
import xml.etree.ElementTree as ET

classes = ['shrinkage', 'blow', 'slag']


# 将训练集图像的名称写在指定文件中,返回文件总数
def write_img_names(img_names, cast_annotations):
    name_list = os.listdir(cast_annotations)
    with open(img_names, 'w')as f:
        for ann in name_list:
            f.write(ann[:-4] + "\n")
    return (len(name_list))


# 将XML中标注信息写在指定文件中
def write_ann(img_names, cast_ann, cast_selected, cast_annotations):
    with open(img_names, 'r')as f:
        txt = f.readlines()
        image_inds = [line.strip() for line in txt]

    with open(cast_ann, 'w') as f:
        for image_ind in image_inds:
            image_path = os.path.join(cast_selected, image_ind + '.png')
            annotation = image_path
            label_path = os.path.join(cast_annotations, image_ind + '.xml')
            root = ET.parse(label_path).getroot()
            objects = root.findall('object')
            for obj in objects:
                difficult = obj.find('difficult').text.strip()
                if (not False) and (int(difficult) == 1):
                    continue
                bbox = obj.find('bndbox')
                class_ind = classes.index(obj.find('name').text.lower().strip())
                xmin = bbox.find('xmin').text.strip()
                xmax = bbox.find('xmax').text.strip()
                ymin = bbox.find('ymin').text.strip()
                ymax = bbox.find('ymax').text.strip()
                annotation += ' ' + ','.join([xmin, ymin, xmax, ymax, str(class_ind)])
            print(annotation)
            f.write(annotation + "\n")


if __name__ == '__main__':
    # 训练集所在路径
    cast_path = "E:\\CAST_PNG\\CAST_Train"
    # 训练集中图像
    cast_selected = os.path.join(cast_path, "CAST_selected")
    # 训练集中标签
    cast_annotations = os.path.join(cast_path, "Annotations")
    # 存放图像名称的文件
    img_names = "../data/filenames/cast_train_name.txt"
    # 存放标签的文件
    cast_ann = "../data/dataset/cast_train.txt"

    if os.path.exists(img_names): os.remove(img_names)
    if os.path.exists(cast_ann): os.remove(cast_ann)

    write_img_names(img_names, cast_annotations)
    write_ann(img_names, cast_ann, cast_selected, cast_annotations)

